# Edge Impulse - OpenMV Object Detection Example
import sensor, image, time,pyb,omv,math,utime,tf,lcd,gc
from pyb import UART,Pin,Timer,Servo
from umotor import UMOTOR
from pid import PID
from ultrasonic import ULTRASONIC
from button import BUTTON
from centroidtracking import CentroidTracking
from pycommon import PYCOMMON
common=PYCOMMON()
objtraking=CentroidTracking(maxDisappeared=5) #目标追踪算法，质心追踪maxDisappeared表示最大丢失5帧代表目标消失
lcd.init(type=2,width=240,height=320)
lcd.set_direction(2)
button=BUTTON()  #声明按键，梦飞openmv只有一个按键，因此直接内部指定了按键
object_s=0
flag_lost=0
max_object=(0,0,0,0,0.0)
pan_servo=Servo(3)    #左右控制PD12
tilt_servo=Servo(4)  #上下控制PD13
pan_start_angle=0.0
tilt_start_angle=-20.0 #朝上更容易看到人脸

pan_angle=pan_servo.angle(pan_start_angle)    #上下运动控制设置一个初始角度
tilt_angle=tilt_servo.angle(tilt_start_angle)   #左右舵运动控制，初始角度
#追踪PID
pan_pid = PID(p=0.2, i=0, imax=90) #脱机运行或者禁用图像传输，使用这个PID
tilt_pid = PID(p=0.2, i=0, imax=90) #脱机运行或者禁用图像传输，使用这个PID
net = None
labels = None
confidence=0.7

try:
    # load the model, alloc the model file on the heap if we have at least 64K free after loading
    labels,net = tf.load_builtin_model('yoloface')
    #net = tf.load("yoloface.tflite",load_to_fb=True)
except Exception as e:
    raise Exception('Failed to load "yoloface", did you copy the .tflite and labels.txt file onto the mass-storage device? (' + str(e) + ')')

print(net)


def face_detect(img):
    max_object=(0,0,0,0,0.0)
    rid=img.width()/img.height()
    xrect=[]
    object_s=0.0
    objects=net.detect_yolo(img,confidence=0.75, anchors=[21,28,34,49,61,77],nms=0.2)#anchors=[9,14,12,17,22,21]#anchors=[21,28,34,49,61,77]
    #if (len(objects) == 0): continue # no detections for this class?
    if len(objects):
        for d in objects :
            #rect=(d.x(), d.y(), d.w(), d.h())
            rect=(int(d.x()/rid), d.y(), int(d.w()/rid), d.h())
            img.draw_rectangle(rect,color=(255,255,0)) #检测为白色框
            img.draw_string(d.x(), d.y()-10, "face %.3f"%(d.output()))
            xrect.append(rect)
    objects=objtraking.update(xrect)
    max_object=common.find_max_object(objects)
    if (max_object):
        object_s=15000/(max_object[2]*2) #计算距离
    return object_s,max_object

#功能： 自动找球控制
#输入： 图像，小球距离,最大面积的小球元组
#输出： 无
def face_traking_auto_control(img):
    global flag_lost
    object_s=0
    object_s,max_blob=face_detect(img)
    if object_s>0:
        pan_error=0
        flag_lost=0
        pan_error = img.width()/2-(max_blob[2]/2+max_blob[0])   #左右控制的偏差
        tilt_error =img.height()/2-(max_blob[3]/2+max_blob[1]) #上下控制的偏差
        pan_output=pan_pid.get_pid(pan_error,1)/2   #计算PID参数
        tilt_output=tilt_pid.get_pid(tilt_error,1)/2  #计算PID参数
        pan_angle=pan_servo.angle()+pan_output
        tilt_angle=tilt_servo.angle()-tilt_output
        if pan_angle>=60:   pan_angle=60
        if pan_angle<=-60:   pan_angle=-60
        if tilt_angle>=60:   tilt_angle=60
        if tilt_angle<=-60:   tilt_angle=-60
        pan_servo.angle(pan_angle)
        tilt_servo.angle(tilt_angle)
    else :
        flag_lost=flag_lost+1
        if flag_lost>5:#连续5帧没有
            flag_lost=0
            pan_angle=pan_servo.angle(pan_start_angle)    #上下运动控制设置一个初始角度
            tilt_angle=tilt_servo.angle(tilt_start_angle)   #左右舵运动控制，初始角度


sensor.reset()                         # Reset and initialize the sensor.
sensor.set_pixformat(sensor.GRAYSCALE)    # Set pixel format to RGB565 (or GRAYSCALE)
sensor.set_framesize(sensor.QVGA)      # Set frame size to QVGA (320x240)
#sensor.set_hmirror(1)
#sensor.set_windowing((200, 200))       # Set 240x240 window.
sensor.set_hmirror(True) #水平镜像，暂时不用
sensor.set_vflip(True)  #垂直镜像，根据摄像头的安装位置调整
sensor.skip_frames(10) # Let new settings take affect.
click_timer=time.ticks() #计时参数
clock = time.clock()
while(True):
    clock.tick()
    img = sensor.snapshot()
    face_traking_auto_control(img)
    print(clock.fps(), "fps", end="\n\n")
    lcd.display(img)
    if button.state():
        click_timer=time.ticks()          #开始计时
        while button.state():  pass       #等待按键抬起
        if time.ticks()-click_timer>2000: #按键时长超过2s
            break                         #循环退出，回到主界面
    else :
        click_timer=time.ticks()#计时更新

